/**
 * Deep4jNetworkWrapper.java created by zhangzhidong 
 * at 上午10:07:50 2017年5月26日
 */
package cn.edu.bjtu.model;

import java.io.File;
import java.io.IOException;

import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.nn.api.Layer;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.api.IterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.cpu.nativecpu.NDArray;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;

import cn.edu.bjtu.configuration.TextCategorizationCNNConfig;
import cn.edu.bjtu.core.Deep4jModelType;
import cn.edu.bjtu.core.LoggerSupport;
import cn.edu.bjtu.core.NeuronNetwork;
import cn.edu.bjtu.model.TextCategorizationCNNModel;

/**
 * 这个里面的方法全是同步执行的，public 方法都是在调用它的线程中执行的,并发控制以及异常控制在{@link TextCategorizationCNNModel}里面控制;
 * 所以NeuronNetwork里面的方法都是抛出异常的
 * 目前来说 ， 是对于 DEEP4J 里面两种模型的包装，，以后模型多了，这个类是可以横向拓展的
 * 如果没有Slave这个类以及它的子类用于封闭不同的网络屏蔽不同网络接口之间接口差异
 * 所以每次次调用Deep4jNetworkWrapper这个类,需要判断当前网络的类型
 * if(类型为MLN)调用MLN的方法
 * if(类型为CG)调用CG的方法
 * 这样代码看起来给人不好的感觉
 * Slave存在的意思根据Java里面的多态使得这些判断可以省略
 * {@link TextCategorizationCNNConfig}
 * {@link Deep4jModelType}
 * @author Alex
 * 
 */
public class Deep4jNetworkWrapper extends LoggerSupport implements NeuronNetwork{
	//使用了这个 salve 类之后。对于Deep4jNetworkWrapper这个类暴露出的方法执行时候不用每次 switch(type)了
	private abstract class Slave extends LoggerSupport implements NeuronNetwork{
		abstract Object get();
	}
	private Slave slave = null;
	private final TextCategorizationCNNModel facade ;
	private final TextCategorizationCNNConfig config;
	public Deep4jNetworkWrapper(TextCategorizationCNNModel facade) {
		this.facade = facade;
		config = facade.getCNNConfig();
		//用了一个 Slave 类，这样只写一次 switch 语句就够了
		switch(config.getDeep4jModelType()){
		case ComputationGraph:
			slave = new ComputationGraphSlave();
			break;
		case MultiLayerNetwork:
			slave  = new MultiLayerNetworkSlave();
			break;
		default:
			throw new IllegalArgumentException("unknow type");
	}
		
	}

	 public INDArray[] output(INDArray... input){
		 return slave.output(input);
	 }	
	@Override
	public void fit(DataSetIterator iterator) throws Exception {
		slave.fit(iterator);
		
	}
	@Override
	public Evaluation evaluate(DataSetIterator iterator) throws Exception {
		return slave.evaluate(iterator);
	}
	private class ComputationGraphSlave extends Slave{
		private ComputationGraph cg = null;
		public ComputationGraphSlave() {
			
		}
		public INDArray[] output(INDArray... input) {
			return cg.output(input);
		}
		public void fit(DataSetIterator iterator) {
			cg.fit(iterator);
		}
		public Evaluation evaluate(DataSetIterator iterator) {
			return cg.evaluate(iterator);
		}
		@Override
		public void saveModel(String path) {
			try {
				ModelSerializer.writeModel(cg, path, true);
			} catch (IOException e) {
				logger.error("Model write failed : {}",e.getMessage(),e);
			}
		}		
		@Override
		public void load() throws IOException {
			cg = ModelSerializer.restoreComputationGraph(new File(config.getNetworkSaveFile()), true);
		}

		@Override
		public void setListeners(IterationListener ... listeners) {
			cg.setListeners(listeners);
			
		}
		
		@Override
		public Layer[] getLayers() {
			return cg.getLayers();
		}
		@Override
		public Deep4jModelType getActualType() {
			return Deep4jModelType.ComputationGraph;
		}		
		@Override
		public void init() {
			cg = new ComputationGraph(facade.getNetworkDesignHandler().handleCGC(facade.getW2v().getLayerSize()));
			cg.init();
		}
		@Override
		public Object get() {
			return cg;
		}
	}
	
	private class MultiLayerNetworkSlave extends Slave{
		private MultiLayerNetwork mln = null;
		public MultiLayerNetworkSlave() {
			
		}
		@Override
		/**
		 * 只能对一个进行分类，传进来是数组也没用
		 * @Param
		 * @Author zhangzhidong
		 */
		public INDArray[] output(INDArray... input) {
			INDArray []res = new NDArray[1];
			res[1] = mln.output(input[0]);
			return res;
		}
		@Override
		public void fit(DataSetIterator iterator) {
			 mln.fit(iterator);
			 
		}
		@Override
		public Evaluation evaluate(DataSetIterator iterator) {
			return mln.evaluate(iterator);
		}		
		@Override
		public void saveModel(String path) throws IOException {
			ModelSerializer.writeModel(mln, path, true);
		}
		
		@Override
		public void load() throws IOException {
			mln = ModelSerializer.restoreMultiLayerNetwork(new File(config.getNetworkSaveFile()), true);
		}
		@Override
		public Layer[] getLayers() {
			return mln.getLayers();
		}
		@Override
		public void setListeners(IterationListener... listeners) {
			mln.setListeners(listeners);
		}
		@Override
		public Deep4jModelType getActualType() {
			return Deep4jModelType.MultiLayerNetwork;
		}

		@Override
		public void init() {
			mln = new MultiLayerNetwork(facade.getNetworkDesignHandler().handleMLN(facade.getW2v().getLayerSize()));
			mln.init();
		}
		@Override
		public Object get() {
			return mln;
		}
		
	}
	@Override
	public void saveModel(String path) throws Exception {
		slave.saveModel(path);
	}

	@Override
	public void load() throws Exception {
		//这个再次检查是有必要的,因为配置文件可能改变
		switch(config.getDeep4jModelType()){
		case ComputationGraph:
			slave = new ComputationGraphSlave();
			break;
		case MultiLayerNetwork:
			slave  = new MultiLayerNetworkSlave();
			break;
		default:
			throw new IllegalArgumentException("unknow type");
	}
		slave.load();
	}

	@Override
	public void setListeners(IterationListener... listeners) {
		slave.setListeners(listeners);
		
	}
	
	@Override
	public Layer[] getLayers() {
		return slave.getLayers();
	}
	
	@Override
	public Deep4jModelType getActualType() {
		return slave.getActualType();
	}
	@Override
	public void init() throws Exception {
	
		slave.init();
	}
	public Object get(){
		return slave.get();
	}
}
